from pprint import pprint
from paddlenlp import Taskflow
import paddle

# paddle.device.set_device("gpu:3")

schema = [{"歌曲名称":["歌手", "所属专辑"]}, "人物"]
ie = Taskflow('information_extraction', schema=schema, model="uie-base")
res = ie("告白气球是周杰伦在专辑周杰伦的床边故事里面的歌曲")
pprint(res)

schema = ["歌曲名称", "人物"]
ie = Taskflow('information_extraction', schema=schema, model="uie-base")
res = ie("告白气球是周杰伦在专辑周杰伦的床边故事里面歌曲")
pprint(res)
